import random
import sys

from utils.client import Client
from utils.router import Router

sys.path.insert(0, './')
sys.path.insert(0, '../')
import numpy as np
import torch.utils.data as data
from PIL import Image
import torch
import os
import torchvision.transforms as transforms
from dataloaders.helper import CutoutPIL
from randaugment import RandAugment
import xml.dom.minidom


class voc2007(data.Dataset):
    def __init__(self, root, data_split, img_size=224, p=1, annFile="", label_mask=None, partial=1 + 1e-6,
                 netpath=None):
        # data_split = train / val
        self.root = root
        self.classnames = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
                           'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa',
                           'train', 'tvmonitor']

        # note 用于cdn数据集
        self.netpath = netpath  # 网络拓扑
        self.net = None
        self.routers = None
        self.clients = None
        self.subscriptions = None
        self.client2idx = None
        if self.netpath:
            self.net, self.routers, self.clients, self.client2idx = self.create_network()
            self.subscriptions = torch.zeros(len(self.classnames), len(self.clients), dtype=int)
            self.assign_content()

        if annFile == "":
            self.annFile = os.path.join(self.root, 'Annotations')
        else:
            raise NotImplementedError

        image_list_file = os.path.join(self.root, 'ImageSets', 'Main', '%s.txt' % data_split)

        with open(image_list_file) as f:
            image_list = f.readlines()
        self.image_list = [a.strip() for a in image_list]

        self.data_split = data_split
        if data_split == 'Train':
            num_examples = len(self.image_list)
            pick_example = int(num_examples * p)
            self.image_list = self.image_list[:pick_example]
        else:
            self.image_list = self.image_list

        train_transform = transforms.Compose([
            # transforms.RandomResizedCrop(img_size)
            transforms.Resize((img_size, img_size)),
            CutoutPIL(cutout_factor=0.5),
            RandAugment(),
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
        ])
        test_transform = transforms.Compose([
            # transforms.CenterCrop(img_size),
            transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
        ])

        if self.data_split == 'trainval':
            self.transform = train_transform
        elif self.data_split == 'test':
            self.transform = test_transform
        else:
            raise ValueError('data split = %s is not supported in Nus Wide' % self.data_split)

        # create the label mask
        self.mask = None
        self.partial = partial
        if data_split == 'trainval' and partial < 1.:
            if label_mask is None:
                rand_tensor = torch.rand(len(self.image_list), len(self.classnames))
                mask = (rand_tensor < partial).long()
                mask = torch.stack([mask], dim=1)
                torch.save(mask, os.path.join(self.root, 'Annotations', 'partial_label_%.2f.pt' % partial))
            else:
                mask = torch.load(os.path.join(self.root, 'Annotations', label_mask))
            self.mask = mask.long()

    def __len__(self):
        return len(self.image_list)

    def __getitem__(self, index):
        img_path = os.path.join(self.root, 'JPEGImages', self.image_list[index] + '.jpg')
        img = Image.open(img_path).convert('RGB')
        ann_path = os.path.join(self.annFile, self.image_list[index] + '.xml')
        label_vector = torch.zeros(20)
        DOMTree = xml.dom.minidom.parse(ann_path)
        root = DOMTree.documentElement
        objects = root.getElementsByTagName('object')
        for obj in objects:
            if (obj.getElementsByTagName('difficult')[0].firstChild.data) == '1':
                continue
            tag = obj.getElementsByTagName('name')[0].firstChild.data.lower()
            label_vector[self.classnames.index(tag)] = 1.0
        targets = label_vector.long()
        target = targets[None,]
        if self.mask is not None:
            masked = - torch.ones((1, len(self.classnames)), dtype=torch.long)
            target = self.mask[index] * target + (1 - self.mask[index]) * masked

        if self.transform is not None:
            img = self.transform(img)

        return img, target

    def create_network(self):
        """
        创建网络
        @return:
        """
        net = {}  # r_name : neighbors
        clients = {}  # name : object
        routers = {}  # name : object
        client2idx = {}  # client_name: idx
        with open(self.netpath) as f:
            is_first = True
            for line in f:
                line = line.strip()
                if is_first:
                    is_first = False
                    nodes = line.split(" ")
                    for node in nodes:
                        if 'R' in node:
                            net[node] = []
                            routers[node] = Router(node)
                        else:
                            clients[node] = Client(node)
                else:
                    nodes = line.split(" ")
                    a = nodes[0]  # 第一个节点
                    b = nodes[1]  # 第二个节点
                    net[b].append(a)  # 第一个节点连接到第二个节点上
            for idx, client in enumerate(clients.keys()):
                client2idx[client] = idx
        return net, routers, clients, client2idx

    def assign_content(self):
        """
        为终端节点分配订阅信息

        @return:
        """
        random.seed(0)
        for name, client in self.clients.items():
            index = random.randint(0, len(self.classnames) - 1)
            self.subscriptions[index][self.client2idx[name]] += 1
            client.set_content(index, self.classnames[index])

    def name(self):
        return 'voc2007'
